package com.shanshan.wheel.rpc.common.registry.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

import javax.annotation.PostConstruct;
import java.net.InetAddress;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @Description zk注册中心工具类
 * @ClassName CuratorUtil
 * @Author 山山大魔王
 * @Date 2022/9/3 14:10
 * @Version 1.0
 */
@Slf4j
@Component
public class CuratorUtil {

    private static final int BASE_SLEEP_TIME = 1000;

    private static final int MAX_RETRIES = 3;

    private static String ZK_ADDRESS;

    private static Integer PORT;

    public static final String ZK_REGISTER_ROOT_PATH = "/my-rpc";

    // 缓存已注册服务的地址
    private static final Map<String, List<String>> SERVICE_ADDRESS_MAP = new ConcurrentHashMap<>();

    // 缓存所有注册地址
    private static final Set<String> REGISTERED_PATH_SET = ConcurrentHashMap.newKeySet();
    // zk客户端操作api
    private static volatile CuratorFramework zkClient;

    @Value("${rpc.zk.address}")
    private String zkAddress;

    @Value("${rpc.zk.port}")
    private Integer port;

    @PostConstruct
    public void getEnvironment(){
        ZK_ADDRESS = this.zkAddress;
        PORT = this.port;
    }

    public CuratorUtil() {
    }

    /**
     * 返回zk客户端（单例）
     * @return
     */
    public static CuratorFramework getClient(){

        if(zkClient == null){
            synchronized (CuratorUtil.class){
                if(zkClient == null){
                    Assert.notNull(ZK_ADDRESS, "未配置 rpc.zk.address");

                    Assert.notNull(PORT, "未配置 rpc.zk.port");

                    RetryPolicy retryPolicy = new ExponentialBackoffRetry(BASE_SLEEP_TIME, MAX_RETRIES);

                    zkClient = CuratorFrameworkFactory.builder()
                            .connectString(ZK_ADDRESS + ":" + PORT)
                            .retryPolicy(retryPolicy)
                            .build();
                    zkClient.start();

                    try {
                        if(!zkClient.blockUntilConnected(30, TimeUnit.SECONDS)){
                            throw new RuntimeException("连接ZK超时！");
                        }
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }
        }

        Assert.notNull(zkClient, "zkClient未创建");
        Assert.state(zkClient.getState() == CuratorFrameworkState.STARTED, "zkClient已关闭");

        return zkClient;
    }

    /**
     * 创建持久节点
     * @param zkClient
     * @param path
     */
    public static void createPersistentNode(CuratorFramework zkClient, String path){
        try {
            //先查看缓存中是否已注册该服务
            if (REGISTERED_PATH_SET.contains(path) || zkClient.checkExists().forPath(path) != null) {
                log.info("节点[{}]已存在", path);
            } else {
                zkClient.create().creatingParentsIfNeeded().forPath(path);
                log.info("节点[{}]创建成功", path);
            }
            REGISTERED_PATH_SET.add(path);
        }catch (Exception e){
            log.error("为路径[{}]创建持久节点失败",path);
        }
    }

    /**
     * 删除该服务端的所有服务
     * @param zkClient zk客户端
     * @param inetAddress 提供服务的主机
     */
    public static void clearRegistry(CuratorFramework zkClient, InetAddress inetAddress){
        log.info("All registered services on the server are cleared:[{}]", REGISTERED_PATH_SET.toString());
        REGISTERED_PATH_SET.stream().parallel().forEach(p -> {
            if(p.endsWith(inetAddress.toString())){
                try {
                    zkClient.delete().forPath(p);
                } catch (Exception exception) {
                    log.error("清除路径 [{}] 的注册表失败" , p);
                }
            }
        });
    }

    /**
     * 获取当前服务地址
     * @param zkClient zk客户端
     * @param rpcServiceName 服务名
     * @return
     */
    public static List<String> getServiceAddress(CuratorFramework zkClient, String rpcServiceName){
        // 先从缓存中取地址
        if(SERVICE_ADDRESS_MAP.containsKey(rpcServiceName)){
            return SERVICE_ADDRESS_MAP.get(rpcServiceName);
        }
        // 缓存中没有则请求zk获取服务地址
        List<String> result = null;

        String servicePath = ZK_REGISTER_ROOT_PATH + "/" + rpcServiceName;

        try {
            result = zkClient.getChildren().forPath(servicePath);
            SERVICE_ADDRESS_MAP.put(rpcServiceName, result);
            // 给该路径注册观察者
            registerWatcher(zkClient, rpcServiceName);
        } catch (Exception exception) {
            log.error("获取路径[{}]的子节点失败",servicePath);
        }
        return result;
    }

    private static void registerWatcher(CuratorFramework zkClient, String rpcServiceName) throws Exception {
        String servicePath = ZK_REGISTER_ROOT_PATH +"/"+rpcServiceName;
        PathChildrenCache cache = new PathChildrenCache(zkClient, servicePath, true);
        cache.getListenable().addListener((curatorFramework , pathChildrenCacheEvent) -> {
            switch (pathChildrenCacheEvent.getType()){
                case CHILD_ADDED:
                    log.info("新注册的服务：[{}]",pathChildrenCacheEvent.getData());
                    break;
                case CHILD_REMOVED:
                    log.info("服务[{}]被移除",pathChildrenCacheEvent.getData());
                    break;
                case CHILD_UPDATED:
                    log.info("服务更新：[{}]",pathChildrenCacheEvent.getData());
                    break;
            }

            List<String> serviceAddresses = curatorFramework.getChildren().forPath(servicePath);
            SERVICE_ADDRESS_MAP.put(rpcServiceName, serviceAddresses);
        });

        cache.start();

    }
}
